#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Generate a 640x384 BMP test image (left black, center white, right red),
then convert it into two plane arrays compatible with Waveshare 7.5" B&C panel.
Outputs:
- docs/test_epd_stripes.bmp (24-bit BMP)
- main/test_image_planes.h (C header with TEST_BW_PLANE and TEST_RY_PLANE arrays)

Semantics:
- Black plane bit: 0 = black, 1 = white
- RY plane bit:    0 = red,  1 = no red
"""
import struct
from pathlib import Path

W = 640
H = 384

root = Path(__file__).resolve().parents[1]
out_bmp = root / 'docs' / 'test_epd_stripes.bmp'
out_hdr = root / 'main' / 'test_image_planes.h'


def make_stripes_image():
    """Return a list of rows, each row is list of (B,G,R) tuples."""
    img = []
    x_left_end = W // 3
    x_right_start = (W * 2) // 3
    for y in range(H):
        row = []
        for x in range(W):
            if x < x_left_end:
                # left: black
                row.append((0, 0, 0))
            elif x >= x_right_start:
                # right: red
                row.append((0, 0, 255))  # BGR order (blue, green, red)
            else:
                # middle: white
                row.append((255, 255, 255))
        img.append(row)
    return img


def write_bmp24(filepath, img):
    """Write 24-bit BMP with bottom-up rows and per-row padding to 4 bytes."""
    height = len(img)
    width = len(img[0])
    bpp = 24
    row_bytes = width * 3
    pad = (4 - (row_bytes % 4)) % 4
    pixel_data_size = (row_bytes + pad) * height

    # BMP headers
    bfType = b'BM'
    bfSize = 14 + 40 + pixel_data_size
    bfReserved1 = 0
    bfReserved2 = 0
    bfOffBits = 14 + 40

    biSize = 40
    biWidth = width
    biHeight = height  # positive for bottom-up
    biPlanes = 1
    biBitCount = bpp
    biCompression = 0
    biSizeImage = pixel_data_size
    biXPelsPerMeter = 0
    biYPelsPerMeter = 0
    biClrUsed = 0
    biClrImportant = 0

    with open(filepath, 'wb') as f:
        # BITMAPFILEHEADER
        f.write(struct.pack('<2sIHHI', bfType, bfSize, bfReserved1, bfReserved2, bfOffBits))
        # BITMAPINFOHEADER
        f.write(struct.pack('<IIIHHIIIIII', biSize, biWidth, biHeight, biPlanes, biBitCount,
                            biCompression, biSizeImage, biXPelsPerMeter, biYPelsPerMeter,
                            biClrUsed, biClrImportant))
        # Pixel data: bottom-up
        for y in reversed(range(height)):
            row = img[y]
            for (b, g, r) in row:
                f.write(struct.pack('<BBB', b, g, r))
            if pad:
                f.write(b'\x00' * pad)


def planes_from_img(img):
    """Convert BGR tuples to two planes according to semantics.
    Returns bytes(bw) and bytes(ry) of length (W*H/8).
    """
    width_bytes = W // 8
    total = width_bytes * H
    bw = bytearray([0xFF] * total)  # 1=white
    ry = bytearray([0xFF] * total)  # 1=no red

    def set_bit(buf, x, y, value):
        byte_index = y * width_bytes + (x // 8)
        mask = 0x80 >> (x & 7)
        if value:
            buf[byte_index] |= mask
        else:
            buf[byte_index] &= (~mask) & 0xFF

    for y in range(H):
        for x in range(W):
            b, g, r = img[y][x]
            if r > 200 and g < 50 and b < 50:
                # red
                set_bit(bw, x, y, True)
                set_bit(ry, x, y, False)
            elif r < 50 and g < 50 and b < 50:
                # black
                set_bit(bw, x, y, False)
                set_bit(ry, x, y, True)
            else:
                # white
                set_bit(bw, x, y, True)
                set_bit(ry, x, y, True)
    return bytes(bw), bytes(ry)


def write_planes_header(filepath, bw, ry):
    with open(filepath, 'w') as f:
        f.write('// Auto-generated by tools/make_epd_bmp_planes.py\n')
        f.write('#pragma once\n\n')
        f.write('#include "DEV_Config.h"\n\n')
        f.write('static const UBYTE TEST_BW_PLANE[%d] = {\n' % len(bw))
        for i in range(0, len(bw), 16):
            chunk = ', '.join('0x%02X' % b for b in bw[i:i+16])
            f.write('  ' + chunk + ',\n')
        f.write('};\n\n')
        f.write('static const UBYTE TEST_RY_PLANE[%d] = {\n' % len(ry))
        for i in range(0, len(ry), 16):
            chunk = ', '.join('0x%02X' % b for b in ry[i:i+16])
            f.write('  ' + chunk + ',\n')
        f.write('};\n')


def main():
    img = make_stripes_image()
    out_bmp.parent.mkdir(parents=True, exist_ok=True)
    out_hdr.parent.mkdir(parents=True, exist_ok=True)
    write_bmp24(out_bmp, img)
    bw, ry = planes_from_img(img)
    write_planes_header(out_hdr, bw, ry)
    print('Generated:', out_bmp)
    print('Generated header:', out_hdr)

if __name__ == '__main__':
    main()
